from annoy import AnnoyIndex
import random
import matplotlib.pyplot as plt
import numpy as np

# f = 40
# t = AnnoyIndex(f, 'angular')
# for i in range(1000):
#     v = [random.gauss(0,1) for z in range(f)]
#     t.add_item(i ,v)
#
# t.build(5)
# t.save('test.ann')
#
# u= AnnoyIndex(f, 'angular')
# u.load('test.ann')
# print(u.get_nns_by_item(0, 10))
# print(u.get_n_trees())

# NSW


class Node:

    def __init__(self, point, neighbors, top_k=3):
        assert len(neighbors) <= top_k, '传入邻居的长度过长'
        self.point = point
        self.neighbors = neighbors


def get_cosine_similarity(point_values):
    dots = point_values.dot(point_values.T)
    norms = np.sqrt(np.sum(point_values**2, axis=1)).reshape(-1, 1)
    return dots / (norms.dot(norms.T))


def naive_arg_topK(matrix, K, axis=0):
    """
    perform topK based on np.argsort
    :param matrix: to be sorted
    :param K: select and sort the top K items
    :param axis: dimension to be sorted.
    :return:
    """
    shape_axis = matrix.shape[axis]
    full_sort = np.argsort(matrix, axis=axis)
    return full_sort.take(np.arange(shape_axis)[shape_axis - K:], axis=axis)


num_points = 10
size = 2
indices = np.arange(num_points)
values = np.random.random((num_points, size))
similarity = get_cosine_similarity(values)
print(similarity)
print(similarity.shape)

indx = naive_arg_topK(similarity, 4, axis=1)
print(indx)
colors = ['b']


